package com.peng.leetcode.array;

import java.util.PriorityQueue;

/**
 * FindMedianSortedArrays
 * <p>
 * 4. 寻找两个正序数组的中位数
 * <p>
 * 给定两个大小为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的中位数。
 * <p>
 * 进阶：你能设计一个时间复杂度为 O(log (m+n)) 的算法解决此问题吗？
 * <p>
 * 链接：https://leetcode-cn.com/problems/median-of-two-sorted-arrays
 *
 * @author: lupeng6
 * @create: 2021/1/1 20:14
 */
public class FindMedianSortedArrays {

    public static void main(String[] args) {
        int[] nums1 = {1, 2}, nums2 = {3, 4};
        System.out.println(new FindMedianSortedArrays().findMedianSortedArrays1(nums1, nums2));
    }

    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int[] arr = new int[nums1.length + nums2.length];
        int index = 0;
        int index1 = 0;
        int index2 = 0;
        while (index1 < nums1.length && index2 < nums2.length) {
            if (nums1[index1] <= nums2[index2]) {
                arr[index] = nums1[index1];
                index1++;
            } else {
                arr[index] = nums2[index2];
                index2++;
            }
            index++;
        }
        while (index1 < nums1.length) {
            arr[index] = nums1[index1];
            index1++;
            index++;
        }
        while (index2 < nums2.length) {
            arr[index] = nums2[index2];
            index2++;
            index++;
        }
        if (arr.length % 2 == 0) {
            return (double) (arr[arr.length / 2 - 1] + arr[arr.length / 2]) / 2;
        } else {
            return arr[arr.length / 2];
        }
    }

    /**
     * 通过堆, 寻找 中位数
     *
     * @author lupeng6
     * @date 2021/1/22 16:45
     */
    public double findMedianSortedArrays1(int[] nums1, int[] nums2) {
        if (nums1.length == 0 && nums2.length == 0) {
            return 0;
        }
        if (nums1.length == 0) {
            if (nums2.length % 2 == 0) {
                return (double) (nums2[nums2.length / 2 - 1] + nums2[nums2.length / 2]) / 2;
            } else {
                return nums2[nums2.length / 2];
            }
        }
        if (nums2.length == 0) {
            if (nums1.length % 2 == 0) {
                return (double) (nums1[nums1.length / 2 - 1] + nums1[nums1.length / 2]) / 2;
            } else {
                return nums1[nums1.length / 2];
            }
        }
        int heapSize = (nums1.length + nums2.length) / 2 + 1;
        PriorityQueue<Integer> minHeap = new PriorityQueue<>(heapSize);
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(heapSize, (a, b) -> b - a);

        if (nums1[0] > nums2[0]) {
            maxHeap.offer(nums2[0]);
            minHeap.offer(nums1[0]);
        } else {
            maxHeap.offer(nums1[0]);
            minHeap.offer(nums2[0]);
        }

        for (int i = 1; i < nums1.length; i++) {
            int num = nums1[i];
            Integer min = minHeap.peek();

            if (min <= num) {
                minHeap.offer(num);
            } else {
                maxHeap.offer(num);
            }

            if (minHeap.size() > maxHeap.size()) {
                maxHeap.offer(minHeap.poll());
            } else if (minHeap.size() < maxHeap.size()) {
                minHeap.offer(maxHeap.poll());
            }
        }

        for (int i = 1; i < nums2.length; i++) {
            int num = nums2[i];
            Integer min = minHeap.peek();

            if (min <= num) {
                minHeap.offer(num);
            } else {
                maxHeap.offer(num);
            }

            if (minHeap.size() > maxHeap.size()) {
                maxHeap.offer(minHeap.poll());
            } else if (minHeap.size() < maxHeap.size()) {
                minHeap.offer(maxHeap.poll());
            }
        }

        if (minHeap.size() == maxHeap.size()) {
            return (double) (minHeap.poll() + maxHeap.poll()) / 2;
        } else if (maxHeap.size() > minHeap.size()) {
            return maxHeap.poll();
        } else {
            return minHeap.poll();
        }
    }

}
